import json

import pandas as pd
from mindsdb_sql_parser.ast import BinaryOperation, Constant, Select
from mindsdb_sql_parser.ast.base import ASTNode

from mindsdb.interfaces.agents.agents_controller import AgentsController
from mindsdb.interfaces.jobs.jobs_controller import JobsController
from mindsdb.interfaces.skills.skills_controller import SkillsController
from mindsdb.interfaces.database.views import ViewController
from mindsdb.interfaces.database.projects import ProjectController
from mindsdb.interfaces.query_context.context_controller import query_context_controller

from mindsdb.api.executor.datahub.datanodes.system_tables import Table


def to_json(obj):
    if obj is None:
        return None
    try:
        return json.dumps(obj)
    except TypeError:
        return obj


def get_project_name(query: ASTNode = None):
    project_name = None
    if (
        isinstance(query, Select)
        and type(query.where) is BinaryOperation
        and query.where.op == "="
        and query.where.args[0].parts == ["project"]
        and isinstance(query.where.args[1], Constant)
    ):
        project_name = query.where.args[1].value
    return project_name


class MdbTable(Table):
    visible: bool = True


class ModelsTable(MdbTable):
    name = "MODELS"
    columns = [
        "NAME",
        "ENGINE",
        "PROJECT",
        "ACTIVE",
        "VERSION",
        "STATUS",
        "ACCURACY",
        "PREDICT",
        "UPDATE_STATUS",
        "MINDSDB_VERSION",
        "ERROR",
        "SELECT_DATA_QUERY",
        "TRAINING_OPTIONS",
        "CURRENT_TRAINING_PHASE",
        "TOTAL_TRAINING_PHASES",
        "TRAINING_PHASE_NAME",
        "TAG",
        "CREATED_AT",
        "TRAINING_TIME",
    ]

    @classmethod
    def get_data(cls, session, inf_schema, **kwargs):
        data = []
        for project_name in inf_schema.get_projects_names():
            project = inf_schema.database_controller.get_project(name=project_name)
            project_models = project.get_models(active=None, with_secrets=session.show_secrets)
            for row in project_models:
                table_name = row["name"]
                table_meta = row["metadata"]

                data.append(
                    [
                        table_name,
                        table_meta["engine"],
                        project_name,
                        table_meta["active"],
                        table_meta["version"],
                        table_meta["status"],
                        table_meta["accuracy"],
                        table_meta["predict"],
                        table_meta["update_status"],
                        table_meta["mindsdb_version"],
                        table_meta["error"],
                        table_meta["select_data_query"],
                        to_json(table_meta["training_options"]),
                        table_meta["current_training_phase"],
                        table_meta["total_training_phases"],
                        table_meta["training_phase_name"],
                        table_meta["label"],
                        row["created_at"],
                        table_meta["training_time"],
                    ]
                )
            # TODO optimise here
            # if target_table is not None and target_table != project_name:
            #     continue

        df = pd.DataFrame(data, columns=cls.columns)
        return df


class DatabasesTable(MdbTable):
    name = "DATABASES"
    columns = ["NAME", "TYPE", "ENGINE", "CONNECTION_DATA"]

    @classmethod
    def get_data(cls, session, inf_schema, **kwargs):
        project = inf_schema.database_controller.get_list(with_secrets=session.show_secrets)
        data = [[x["name"], x["type"], x["engine"], to_json(x.get("connection_data"))] for x in project]

        df = pd.DataFrame(data, columns=cls.columns)
        return df


class MLEnginesTable(MdbTable):
    name = "ML_ENGINES"
    columns = ["NAME", "HANDLER", "CONNECTION_DATA"]

    @classmethod
    def get_data(cls, session, inf_schema, **kwargs):
        integrations = inf_schema.integration_controller.get_all(show_secrets=session.show_secrets)
        ml_integrations = {key: val for key, val in integrations.items() if val["type"] == "ml"}

        data = []
        for _key, val in ml_integrations.items():
            data.append([val["name"], val.get("engine"), to_json(val.get("connection_data"))])

        df = pd.DataFrame(data, columns=cls.columns)
        return df


class HandlersTable(MdbTable):
    name = "HANDLERS"
    columns = [
        "NAME",
        "TYPE",
        "TITLE",
        "DESCRIPTION",
        "VERSION",
        "CONNECTION_ARGS",
        "IMPORT_SUCCESS",
        "IMPORT_ERROR",
    ]

    @classmethod
    def get_data(cls, inf_schema, **kwargs):
        handlers = inf_schema.integration_controller.get_handlers_import_status()

        data = []
        for _key, val in handlers.items():
            connection_args = val.get("connection_args")
            if connection_args is not None:
                connection_args = to_json(connection_args)
            import_success = val.get("import", {}).get("success")
            import_error = val.get("import", {}).get("error_message")
            data.append(
                [
                    val["name"],
                    val.get("type"),
                    val.get("title"),
                    val.get("description"),
                    val.get("version"),
                    connection_args,
                    import_success,
                    import_error,
                ]
            )

        df = pd.DataFrame(data, columns=cls.columns)
        return df


class JobsTable(MdbTable):
    name = "JOBS"
    columns = [
        "NAME",
        "PROJECT",
        "START_AT",
        "END_AT",
        "NEXT_RUN_AT",
        "SCHEDULE_STR",
        "QUERY",
        "IF_QUERY",
        "VARIABLES",
    ]

    @classmethod
    def get_data(cls, query: ASTNode = None, **kwargs):
        jobs_controller = JobsController()

        project_name = None
        if (
            isinstance(query, Select)
            and type(query.where) is BinaryOperation
            and query.where.op == "="
            and query.where.args[0].parts == ["project"]
            and isinstance(query.where.args[1], Constant)
        ):
            project_name = query.where.args[1].value

        data = jobs_controller.get_list(project_name)

        columns = cls.columns
        columns_lower = [col.lower() for col in columns]

        # to list of lists
        data = [[row[k] for k in columns_lower] for row in data]

        return pd.DataFrame(data, columns=columns)


class TriggersTable(MdbTable):
    name = "TRIGGERS"
    columns = [
        "TRIGGER_CATALOG",
        "TRIGGER_SCHEMA",
        "TRIGGER_NAME",
        "EVENT_MANIPULATION",
        "EVENT_OBJECT_CATALOG",
        "EVENT_OBJECT_SCHEMA",
        "EVENT_OBJECT_TABLE",
        "ACTION_ORDER",
        "ACTION_CONDITION",
        "ACTION_STATEMENT",
        "ACTION_ORIENTATION",
        "ACTION_TIMING",
        "ACTION_REFERENCE_OLD_TABLE",
        "ACTION_REFERENCE_NEW_TABLE",
        "ACTION_REFERENCE_OLD_ROW",
        "ACTION_REFERENCE_NEW_ROW",
        "CREATED",
        "SQL_MODE",
        "DEFINER",
        "CHARACTER_SET_CLIENT",
        "COLLATION_CONNECTION",
        "DATABASE_COLLATION",
    ]

    mindsdb_columns = ["NAME", "PROJECT", "DATABASE", "TABLE", "QUERY", "LAST_ERROR"]

    @classmethod
    def get_data(cls, query: ASTNode = None, inf_schema=None, **kwargs):
        from mindsdb.interfaces.triggers.triggers_controller import TriggersController

        triggers_controller = TriggersController()

        project_name = None
        if (
            isinstance(query, Select)
            and type(query.where) is BinaryOperation
            and query.where.op == "="
            and query.where.args[0].parts == ["project"]
            and isinstance(query.where.args[1], Constant)
        ):
            project_name = query.where.args[1].value

        data = triggers_controller.get_list(project_name)

        columns = cls.mindsdb_columns
        if inf_schema.session.api_type == "sql":
            columns = columns + cls.columns
        columns_lower = [col.lower() for col in columns]

        # to list of lists
        data = [[row.get(k) for k in columns_lower] for row in data]

        return pd.DataFrame(data, columns=columns)


class ChatbotsTable(MdbTable):
    name = "CHATBOTS"
    columns = [
        "NAME",
        "PROJECT",
        "DATABASE",
        "MODEL_NAME",
        "PARAMS",
        "IS_RUNNING",
        "LAST_ERROR",
        "WEBHOOK_TOKEN",
    ]

    @classmethod
    def get_data(cls, query: ASTNode = None, **kwargs):
        from mindsdb.interfaces.chatbot.chatbot_controller import ChatBotController

        chatbot_controller = ChatBotController()

        project_name = None
        if (
            isinstance(query, Select)
            and type(query.where) is BinaryOperation
            and query.where.op == "="
            and query.where.args[0].parts == ["project"]
            and isinstance(query.where.args[1], Constant)
        ):
            project_name = query.where.args[1].value

        chatbot_data = chatbot_controller.get_chatbots(project_name=project_name)

        columns = cls.columns
        columns_lower = [col.lower() for col in columns]

        # to list of lists
        data = []
        for row in chatbot_data:
            row["params"] = to_json(row["params"])
            data.append([row[k] for k in columns_lower])

        return pd.DataFrame(data, columns=columns)


class KBTable(MdbTable):
    name = "KNOWLEDGE_BASES"
    columns = [
        "NAME",
        "PROJECT",
        "EMBEDDING_MODEL",
        "RERANKING_MODEL",
        "STORAGE",
        "METADATA_COLUMNS",
        "CONTENT_COLUMNS",
        "ID_COLUMN",
        "PARAMS",
        "INSERT_STARTED_AT",
        "INSERT_FINISHED_AT",
        "PROCESSED_ROWS",
        "ERROR",
        "QUERY_ID",
    ]

    @classmethod
    def get_data(cls, query: ASTNode = None, inf_schema=None, **kwargs):
        project_name = get_project_name(query)

        from mindsdb.interfaces.knowledge_base.controller import KnowledgeBaseController

        controller = KnowledgeBaseController(inf_schema.session)
        kb_list = controller.list(project_name)

        # shouldn't be a lot of queries, we can fetch them all
        queries_data = {item["id"]: item for item in query_context_controller.list_queries()}

        data = []

        for kb in kb_list:
            query_item = {}
            query_id = kb["query_id"]
            if query_id is not None:
                if query_id in queries_data:
                    query_item = queries_data.get(query_id)
                else:
                    query_id = None

            data.append(
                (
                    kb["name"],
                    kb["project_name"],
                    to_json(kb["embedding_model"]),
                    to_json(kb["reranking_model"]),
                    kb["vector_database"] + "." + kb["vector_database_table"],
                    to_json(kb["metadata_columns"]),
                    to_json(kb["content_columns"]),
                    kb["id_column"],
                    to_json(kb["params"]),
                    query_item.get("started_at"),
                    query_item.get("finished_at"),
                    query_item.get("processed_rows"),
                    query_item.get("error"),
                    query_id,
                )
            )

        return pd.DataFrame(data, columns=cls.columns)


class SkillsTable(MdbTable):
    name = "SKILLS"
    columns = ["NAME", "PROJECT", "TYPE", "PARAMS"]

    @classmethod
    def get_data(cls, query: ASTNode = None, **kwargs):
        skills_controller = SkillsController()

        project_name = get_project_name(query)

        all_skills = skills_controller.get_skills(project_name)

        project_controller = ProjectController()
        project_names = {p.id: p.name for p in project_controller.get_list()}

        # NAME, PROJECT, TYPE, PARAMS
        data = [(s.name, project_names[s.project_id], s.type, s.params) for s in all_skills]
        return pd.DataFrame(data, columns=cls.columns)


class AgentsTable(MdbTable):
    name = "AGENTS"
    columns = ["NAME", "PROJECT", "MODEL_NAME", "SKILLS", "PARAMS"]

    @classmethod
    def get_data(cls, query: ASTNode = None, inf_schema=None, **kwargs):
        agents_controller = AgentsController()

        project_name = get_project_name(query)
        all_agents = agents_controller.get_agents(project_name)

        project_controller = ProjectController()
        project_names = {i.id: i.name for i in project_controller.get_list()}

        # NAME, PROJECT, MODEL, SKILLS, PARAMS
        data = [
            (
                a.name,
                project_names[a.project_id],
                a.model_name,
                [rel.skill.name for rel in a.skills_relationships],
                to_json(a.params),
            )
            for a in all_agents
        ]
        return pd.DataFrame(data, columns=cls.columns)


class ViewsTable(MdbTable):
    name = "VIEWS"
    columns = ["NAME", "PROJECT", "QUERY"]

    @classmethod
    def get_data(cls, query: ASTNode = None, **kwargs):
        project_name = get_project_name(query)

        data = ViewController().list(project_name)

        columns_lower = [col.lower() for col in cls.columns]

        # to list of lists
        data = [[row[k] for k in columns_lower] for row in data]

        return pd.DataFrame(data, columns=cls.columns)


class QueriesTable(MdbTable):
    name = "QUERIES"
    columns = [
        "ID",
        "STARTED_AT",
        "FINISHED_AT",
        "PROCESSED_ROWS",
        "ERROR",
        "SQL",
        "DATABASE",
        "PARAMETERS",
        "CONTEXT",
        "UPDATED_AT",
    ]

    @classmethod
    def get_data(cls, **kwargs):
        """
        Returns all queries in progres or recently completed
        Only queries marked as is_resumable by planner are stored in this table
        :param kwargs:
        :return:
        """

        data = query_context_controller.list_queries()
        columns_lower = [col.lower() for col in cls.columns]

        data = [[row[k] for k in columns_lower] for row in data]

        return pd.DataFrame(data, columns=cls.columns)
